package org.budo.ocr.service.tess4j;

import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.net.URL;
import java.util.Collections;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import javax.imageio.ImageIO;

import org.budo.ocr.service.api.BudoOcrService;
import org.budo.ocr.service.api.OcrRequest;
import org.budo.ocr.service.api.OcrResponse;
import org.budo.ocr.service.impl.ByteArrayOcrRequest;
import org.budo.ocr.service.impl.DefaultOcrResponse;
import org.budo.support.java.awt.robot.BudoImage;
import org.budo.support.java.awt.robot.BudoImageImpl;
import org.budo.support.java.net.util.JavaNetUrlUtil;
import org.budo.support.lang.util.FileUtil;
import org.budo.support.lang.util.IoUtil;

import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import net.sourceforge.tess4j.ITesseract;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.util.ImageHelper;

/**
 * @author limingwei
 */
@Getter
@Setter
@Slf4j
public class Tess4jOcrService implements BudoOcrService {
    private static final String TMP_DIR = System.getProperty("java.io.tmpdir") + File.separator + "budo-ocr-service-tess4j-tmpdir";

    private static final Map<String, Tesseract> _TESSERACT_MAP = new ConcurrentHashMap<String, Tesseract>();

    private String dataPath = "tessdata";

    static {
        log.info("#44 delete " + TMP_DIR);
        FileUtil.deleteIfExists(TMP_DIR);
    }

    protected ITesseract getTesseract(OcrRequest request) {
        String language = this.getLanguage(request);

        Tesseract tesseract = _TESSERACT_MAP.get(language);
        if (null != tesseract) {
            return tesseract;
        }

        synchronized (Tess4jOcrService._TESSERACT_MAP) {
            tesseract = _TESSERACT_MAP.get(language);
            if (null != tesseract) {
                return tesseract;
            }

            Tesseract newTesseract = this.newTesseract(request);

            _TESSERACT_MAP.put(language, newTesseract);
            return newTesseract;
        }
    }

    protected Tesseract newTesseract(OcrRequest request) {
        String language = this.getLanguage(request);

        String filePath = this.filePath(this.getDataPath(), language);
        Tesseract tesseractNew = new Tesseract();
        tesseractNew.setDatapath(filePath);
        tesseractNew.setLanguage(language);
        return tesseractNew;
    }

    private String getLanguage(OcrRequest request) {
        String language;
        if (request instanceof Tess4jOcrRequest) {
            language = ((Tess4jOcrRequest) request).getLanguage();
        } else {
            language = "chi_sim";
        }
        return language;
    }

    @Override
    public OcrResponse recognition(OcrRequest request) {
        try {
            byte[] byteArray = ((ByteArrayOcrRequest)request).getByteArray();

            // 右侧消息截取核心部分
            // byteArray = this.subRightMessage(byteArray);

            BufferedImage image = ImageIO.read(new ByteArrayInputStream(byteArray));
            if (null == image) {
                log.error("#65 image is null, request=" + request);
                return null;
            }

            image = this.convertImage(image); // 对图片进行处理 - 提高识别度

            String word = this.getTesseract(request).doOCR(image);
            return new DefaultOcrResponse(Collections.singletonList(word));
        } catch (Throwable e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 截取消息图片的文字部分
     */
    protected byte[] subRightMessage(byte[] byteArray) {
        BudoImage budoImage = new BudoImageImpl(new ByteArrayInputStream(byteArray));
        int width = budoImage.width();
        int y = budoImage.height() / 2;

        // 聊天背景灰色
        int _BACK_GREY = -657931;

        boolean in = false;
        int from = 0;
        for (int x = 50; x < width; x++) { // 左边50个像素是头像
            int rgb = budoImage.rgb(x, y);
            if (_BACK_GREY != rgb && !in) {
                in = true;
                from = x;
                continue;
            }

            if (_BACK_GREY == rgb && in) {
                int x2 = from + 5;
                int w = x - from - 15;
                if (w <= 0) {
                    log.debug("#93 subRightMessage continue, w=" + w + ", x=" + x + ", from=" + from);
                    continue;
                }

                BudoImage sub = budoImage.subImage(x2, 0, w, budoImage.height());
                byte[] png = sub.toPng();

                return png; // 返回截取后的部分
            }
        }

        return byteArray; // 返回本身
    }

    // 对图片进行处理 - 提高识别度
    private BufferedImage convertImage(BufferedImage image) throws Exception {
        // 图像转换成灰度的简单方法 - 黑白处理
        image = ImageHelper.convertImageToGrayscale(image);

        // 图像缩放 - 放大n倍图像
        int targetWidth = image.getWidth() * 3;
        int targetHeight = image.getHeight() * 3;

        return ImageHelper.getScaledInstance(image, targetWidth, targetHeight);
    }

    private String filePath(String dataPath, String language) {
        String fileName = language + ".traineddata";
        File to = new File(TMP_DIR, fileName); // 临时文件

        if (to.exists() && to.length() > 0) {
            return FileUtil.getCanonicalPath(to.getParentFile());
        }

        URL from = this.getClass().getClassLoader().getResource(dataPath + "/" + fileName);

        String protocol = from.getProtocol();
        if ("file".equals(protocol)) {
            String resourceFile = from.getFile();
            File dir = new File(resourceFile).getParentFile();
            String canonicalPath = FileUtil.getCanonicalPath(dir);

            log.info("#47 dataPath=" + dataPath + ", protocol=" + protocol + ", language=" + language + ", canonicalPath=" + canonicalPath);
            return canonicalPath;
        }

        IoUtil.write(JavaNetUrlUtil.openStream(from), to); // 写临时文件

        String canonicalPath = FileUtil.getCanonicalPath(to.getParentFile());

        log.info("#62 dataPath=" + dataPath + ", protocol=" + protocol + ", language=" + language + ", canonicalPath=" + canonicalPath);
        return canonicalPath;
    }
}